import os

from qgis.core import (QgsProcessingAlgorithm, 
                    #    QgsProcessingParameterFolder,
                       QgsProcessingParameterFile,
                       QgsProcessingOutputString,
                       QgsProcessingException)
from qgis.PyQt.QtCore import QCoreApplication
from qgis.core import (QgsProject, QgsVectorLayer, QgsField, QgsPointXY,
                      QgsMarkerSymbol, QgsLineSymbol,
                      QgsFeature, QgsGeometry, QgsFields, 
                      QgsCoordinateReferenceSystem,
                      QgsTextAnnotation,
                      QgsRendererCategory, QgsCategorizedSymbolRenderer)
from PyQt5.QtCore import QVariant
from PyQt5.QtGui import (QColor, QTextDocument)

from qgis.gui import QgsMapCanvasAnnotationItem

COLORS = [
    '#e6194b', '#3cb44b', '#ffe119', '#4363d8', '#f58231',
    '#911eb4', '#46f0f0', '#f032e6', '#bcf60c', '#fabebe',
    '#008080', '#e6beff', '#9a6324', '#fffac8', '#800000',
    '#aaffc3', '#808000', '#ffd8b1', '#000075', '#808080'
]

class SBEB_FDEB_RenderingTool(QgsProcessingAlgorithm):
    """
    This algorithm loads and renders SBEB and FDEB algorithm outputs in QGIS.
    """

    def initAlgorithm(self, config=None):
        self.addParameter(
            QgsProcessingParameterFile(
                'DATAFOLDER',
                self.tr('Input data folder'),
                behavior=QgsProcessingParameterFile.Folder,
                defaultValue=None
            )
        )

        self.addOutput(
            QgsProcessingOutputString(
                'OUTPUT',
                self.tr('Data folder path')
            )
        )

    def processAlgorithm(self, parameters, context, feedback):
        data_folder = self.parameterAsString(parameters, 'DATAFOLDER', context)
        
        if not data_folder:
            raise QgsProcessingException(self.invalidSourceError(parameters, 'DATAFOLDER'))
        
        feedback.pushInfo(f"Loading data from: {data_folder}")
        
        try:
            node_coords = self.load_nodes_layer(data_folder, feedback)
            self.load_edges_layer(data_folder, node_coords, feedback)
            self.load_clustered_edges_layer(data_folder, feedback)
            self.load_skeleton_edges_layer(data_folder, feedback)
            self.load_warped_edges_layer(data_folder, feedback)
            self.load_fdeb_edges_layer(data_folder, feedback)
            feedback.pushInfo("All layers loaded successfully!")
            
        except Exception as e:
            feedback.reportError(f"Error during processing: {str(e)}")
            raise QgsProcessingException(str(e))
        
        return {'OUTPUT': data_folder}

    def name(self):
        return 'sbeb_fdeb_rendering'

    def displayName(self):
        return self.tr('SBEB/FDEB Rendering Tool')

    def group(self):
        return self.tr('Custom Scripts')

    def groupId(self):
        return 'customscripts'

    def tr(self, string):
        return QCoreApplication.translate('Processing', string)

    def createInstance(self):
        return SBEB_FDEB_RenderingTool()


    def load_nodes_layer(self, data_folder, feedback=None):
        nodes_path = os.path.join(data_folder, 'nodes.csv')
        
        if feedback:
            feedback.pushInfo(f"Loading nodes from {nodes_path}")
        
        layer = QgsVectorLayer('Point?crs=EPSG:4326', 'Nodes', 'memory')
        pr = layer.dataProvider()

        pr.addAttributes([
            QgsField('id', QVariant.Int),
            QgsField('lon', QVariant.Double),
            QgsField('lat', QVariant.Double)
        ])
        layer.updateFields()

        node_coords = {}  # id -> QgsPointXY

        with open(nodes_path, 'r') as f:
            next(f)  # 跳过 header
            for line in f:
                parts = line.strip().split(',')
                if len(parts) < 3:
                    continue
                node_id, lon, lat = int(parts[0]), float(parts[1]), float(parts[2])
                point = QgsPointXY(lon, lat)
                feat = QgsFeature()
                feat.setGeometry(QgsGeometry.fromPointXY(point))
                feat.setAttributes([node_id, lon, lat])
                pr.addFeature(feat)

                node_coords[node_id] = point

        layer.updateExtents()
        symbol = QgsMarkerSymbol.createSimple({'name': 'circle', 'color': 'gray', 'size': '2'})
        layer.renderer().setSymbol(symbol)

        QgsProject.instance().addMapLayer(layer)
        print(f'Loaded nodes from {nodes_path} into QGIS.')

        return node_coords

    def load_edges_layer(self, data_folder, node_coords, feedback=None):
        edges_path = os.path.join(data_folder, 'edges.csv')
        
        if feedback:
            feedback.pushInfo(f"Loading edges from {edges_path}")
        
        edges_path = os.path.join(data_folder, 'edges.csv')

        layer = QgsVectorLayer('LineString?crs=EPSG:4326', 'Edges', 'memory')
        pr = layer.dataProvider()

        pr.addAttributes([
            QgsField('from', QVariant.Int),
            QgsField('to', QVariant.Int)
        ])
        layer.updateFields()

        with open(edges_path, 'r') as f:
            header = next(f).strip().split(',')
            has_weight = (len(header) == 3 and header[2].lower() == 'weight')

            for line in f:
                parts = line.strip().split(',')
                if (has_weight and len(parts) < 3) or (not has_weight and len(parts) < 2):
                    continue

                from_id = int(parts[0])
                to_id = int(parts[1])

                # 确保节点存在
                if from_id not in node_coords or to_id not in node_coords:
                    # 也可以打印警告
                    print(f"Warning: Edge from {from_id} to {to_id} skipped because node not found.")
                    continue

                points = [node_coords[from_id], node_coords[to_id]]
                feat = QgsFeature()
                feat.setGeometry(QgsGeometry.fromPolylineXY(points))
                feat.setAttributes([from_id, to_id])
                pr.addFeature(feat)

        layer.updateExtents()

        # 设置简单线符号，比如灰色细线
        symbol = layer.renderer().symbol()
        symbol.setColor(QColor('gray'))
        symbol.setWidth(0.5)

        QgsProject.instance().addMapLayer(layer)
        print(f'Loaded edges from {edges_path} into QGIS.')

    def load_clustered_edges_layer(self, data_folder, feedback=None):
        clustered_path = os.path.join(data_folder, 'gen_clustered_edges.csv')
        
        if feedback:
            feedback.pushInfo(f"Loading clustered edges from {clustered_path}")
        
        # 你的原始实现...
        # ... (保持你的load_clustered_edges_layer函数内容不变)
        clustered_edges_path = os.path.join(data_folder, 'gen_clustered_edges.csv')
        
        # 检查文件是否存在
        if not os.path.exists(clustered_edges_path):
            print(f"Clustered edges file not found at {clustered_edges_path}")
            return
        
        # 获取节点图层
        nodes_layer = QgsProject.instance().mapLayersByName("Nodes")[0]
        if not nodes_layer:
            print("Nodes layer not found. Please load nodes first.")
            return
        
        # 创建节点ID到坐标的映射字典
        node_id_to_point = {}
        for feature in nodes_layer.getFeatures():
            node_id = feature['id']
            geom = feature.geometry()
            point = geom.asPoint()
            node_id_to_point[node_id] = point
        
        # 创建内存图层（线）
        layer = QgsVectorLayer('LineString?crs=EPSG:4326', 'Clustered Edges', 'memory')
        pr = layer.dataProvider()
        
        # 定义字段结构
        pr.addAttributes([
            QgsField('cluster', QVariant.Int),
            QgsField('from', QVariant.Int),
            QgsField('to', QVariant.Int)
        ])
        layer.updateFields()
        
        # 读取 CSV 添加线要素
        with open(clustered_edges_path, 'r') as f:
            headers = next(f).strip().split(',')
            
            for line in f:
                parts = line.strip().split(',')
                if len(parts) < 3:
                    continue
                    
                cluster = int(parts[0])
                from_id = int(parts[1])
                to_id = int(parts[2])
                
                # 获取起点和终点坐标
                from_point = node_id_to_point.get(from_id)
                to_point = node_id_to_point.get(to_id)
                
                if from_point and to_point:
                    # 创建线要素
                    line_geom = QgsGeometry.fromPolylineXY([from_point, to_point])
                    feat = QgsFeature()
                    feat.setGeometry(line_geom)
                    feat.setAttributes([cluster, from_id, to_id])
                    pr.addFeature(feat)
        
        layer.updateExtents()
        
        # 设置分类渲染
        categories = []
        
        # 获取唯一的cluster值
        unique_clusters = set()
        for feature in layer.getFeatures():
            unique_clusters.add(feature['cluster'])
        
        # 为每个cluster创建样式
        for cluster in sorted(unique_clusters):
            color_idx = cluster % len(COLORS)
            symbol = QgsLineSymbol.createSimple({
                'line_color': COLORS[color_idx],
                'line_width': '0.5'
            })
            category = QgsRendererCategory(
                cluster,
                symbol,
                str(cluster)
            )
            categories.append(category)
        
        # 应用分类渲染器
        renderer = QgsCategorizedSymbolRenderer('cluster', categories)
        layer.setRenderer(renderer)
        
        # 触发重绘
        layer.triggerRepaint()
        
        QgsProject.instance().addMapLayer(layer)
        print(f'Loaded clustered edges from {clustered_edges_path} into QGIS.')

        #############################

    def load_skeleton_edges_layer(self, data_folder, feedback=None):
        skeleton_path = os.path.join(data_folder, 'gen_skeleton_edges.csv')
        
        if feedback:
            feedback.pushInfo(f"Loading skeleton edges from {skeleton_path}")
        
        skeleton_edges_path = os.path.join(data_folder, 'gen_skeleton_edges.csv')
        
        # 检查文件是否存在
        if not os.path.exists(skeleton_edges_path):
            print(f"Skeleton edges file not found at {skeleton_edges_path}")
            return
        
        # 创建内存图层（线）
        layer = QgsVectorLayer('LineString?crs=EPSG:4326', 'Skeleton Edges', 'memory')
        pr = layer.dataProvider()
        
        # 定义字段结构
        pr.addAttributes([
            QgsField('cluster_id', QVariant.Int),
            QgsField('x1', QVariant.Double),
            QgsField('y1', QVariant.Double),
            QgsField('x2', QVariant.Double),
            QgsField('y2', QVariant.Double)
        ])
        layer.updateFields()
        
        # 读取 CSV 添加线要素
        with open(skeleton_edges_path, 'r') as f:
            headers = next(f).strip().split(',')
            
            for line in f:
                parts = line.strip().split(',')
                if len(parts) < 5:
                    continue
                    
                cluster_id = int(parts[0])
                x1, y1, x2, y2 = map(float, parts[1:5])
                
                # 创建线要素
                line_geom = QgsGeometry.fromPolylineXY([
                    QgsPointXY(x1, y1),
                    QgsPointXY(x2, y2)
                ])
                feat = QgsFeature()
                feat.setGeometry(line_geom)
                feat.setAttributes([cluster_id, x1, y1, x2, y2])
                pr.addFeature(feat)
        
        layer.updateExtents()
        
        # 设置分类渲染（使用与clustered edges相同的颜色方案）
        categories = []
        
        # 获取唯一的cluster_id值并按升序排序
        unique_clusters = sorted({feat['cluster_id'] for feat in layer.getFeatures()})
        
        # 为每个cluster创建样式
        for cluster in unique_clusters:
            # 使用模运算确保颜色循环使用（与之前一致）
            color = COLORS[cluster % len(COLORS)]
            
            # 创建符号（可以调整线宽使其更突出）
            symbol = QgsLineSymbol.createSimple({
                'line_color': color,
                'line_width': '1.0',  # 比普通边稍粗
                'line_style': 'solid'
            })
            
            # 创建分类
            category = QgsRendererCategory(
                cluster,
                symbol,
                f'Skeleton Cluster {cluster}'
            )
            categories.append(category)
        
        # 应用分类渲染器
        renderer = QgsCategorizedSymbolRenderer('cluster_id', categories)
        layer.setRenderer(renderer)
        
        # 触发重绘
        layer.triggerRepaint()
        
        QgsProject.instance().addMapLayer(layer)
        print(f'Loaded skeleton edges from {skeleton_edges_path} into QGIS.')


    def load_warped_edges_layer(self, data_folder, feedback=None):
        warped_path = os.path.join(data_folder, 'gen_warped_edges.csv')
        
        if feedback:
            feedback.pushInfo(f"Loading warped edges from {warped_path}")
        
        # 你的原始实现...
        # ... (保持你的load_warped_edges_layer函数内容不变)
        warped_edges_path = os.path.join(data_folder, 'gen_warped_edges.csv')
        
        # 检查文件是否存在
        if not os.path.exists(warped_edges_path):
            print(f"Warped edges file not found at {warped_edges_path}")
            return
        
        # 创建内存图层（线）
        layer = QgsVectorLayer('LineString?crs=EPSG:4326', 'Warped Edges', 'memory')
        pr = layer.dataProvider()
        
        # 定义字段结构
        pr.addAttributes([
            QgsField('edge_id', QVariant.Int),
            QgsField('cluster', QVariant.Int)
        ])
        layer.updateFields()
        
        # 先读取所有数据并按edge_id和point_index排序
        edges_data = {}
        with open(warped_edges_path, 'r') as f:
            headers = next(f).strip().split(',')
            
            for line in f:
                parts = line.strip().split(',')
                if len(parts) < 5:
                    continue
                    
                edge_id = int(parts[0])
                point_index = int(parts[1])
                lon = float(parts[2])
                lat = float(parts[3])
                cluster = int(parts[4])
                
                if edge_id not in edges_data:
                    edges_data[edge_id] = {
                        'points': [],
                        'cluster': cluster
                    }
                
                edges_data[edge_id]['points'].append((point_index, QgsPointXY(lon, lat)))
        
        # 为每个edge_id创建线要素
        for edge_id, data in edges_data.items():
            # 按point_index排序点
            sorted_points = sorted(data['points'], key=lambda x: x[0])
            points = [point for idx, point in sorted_points]
            
            # 创建线要素
            line_geom = QgsGeometry.fromPolylineXY(points)
            feat = QgsFeature()
            feat.setGeometry(line_geom)
            feat.setAttributes([edge_id, data['cluster']])
            pr.addFeature(feat)
        
        layer.updateExtents()
        
        # 设置分类渲染（使用与之前相同的颜色方案）
        categories = []
        
        # 获取唯一的cluster值并按升序排序
        unique_clusters = sorted({feat['cluster'] for feat in layer.getFeatures()})
        
        # 为每个cluster创建样式
        for cluster in unique_clusters:
            # 使用模运算确保颜色循环使用（与之前一致）
            color = COLORS[cluster % len(COLORS)]
            
            # 创建符号（可以调整线宽或样式以区分不同类型的边）
            symbol = QgsLineSymbol.createSimple({
                'line_color': color,
                'line_width': '0.8',  # 中等线宽
                'line_style': 'solid'
            })
            
            # 创建分类
            category = QgsRendererCategory(
                cluster,
                symbol,
                f'Warped Cluster {cluster}'
            )
            categories.append(category)
        
        # 应用分类渲染器
        renderer = QgsCategorizedSymbolRenderer('cluster', categories)
        layer.setRenderer(renderer)
        
        # 触发重绘
        layer.triggerRepaint()
        
        QgsProject.instance().addMapLayer(layer)
        print(f'Loaded warped edges from {warped_edges_path} into QGIS.')

    def load_fdeb_edges_layer(self, data_folder, feedback=None):

        fdeb_path = os.path.join(data_folder, 'gen_fdeb_edges.csv')
        
        if feedback:
            feedback.pushInfo(f"Loading FDEB edges from {fdeb_path}")
        
        # 你的原始实现...
        # ... (保持你的load_fdeb_edges_layer函数内容不变)
        fdeb_edges_path = os.path.join(data_folder, 'gen_fdeb_edges.csv')
        
        # 检查文件是否存在
        if not os.path.exists(fdeb_edges_path):
            print(f"FDEB edges file not found at {fdeb_edges_path}")
            return
        
        # 创建内存图层（线）
        layer = QgsVectorLayer('LineString?crs=EPSG:4326', 'FDEB Edges', 'memory')
        pr = layer.dataProvider()
        
        # 定义字段结构
        pr.addAttributes([
            QgsField('edge_id', QVariant.Int)
        ])
        layer.updateFields()
        
        # 先读取所有数据并按edge_id和point_index排序
        edges_data = {}
        with open(fdeb_edges_path, 'r') as f:
            headers = next(f).strip().split(',')
            
            for line in f:
                parts = line.strip().split(',')
                if len(parts) < 4:
                    continue
                    
                edge_id = int(parts[0])
                point_index = int(parts[1])
                lon = float(parts[2])
                lat = float(parts[3])
                
                if edge_id not in edges_data:
                    edges_data[edge_id] = []
                
                edges_data[edge_id].append((point_index, QgsPointXY(lon, lat)))
        
        # 为每个edge_id创建线要素
        for edge_id, points_data in edges_data.items():
            # 按point_index排序点
            sorted_points = sorted(points_data, key=lambda x: x[0])
            points = [point for idx, point in sorted_points]
            
            # 至少需要2个点才能形成线
            if len(points) < 2:
                continue
                
            # 创建线要素
            line_geom = QgsGeometry.fromPolylineXY(points)
            feat = QgsFeature()
            feat.setGeometry(line_geom)
            feat.setAttributes([edge_id])
            pr.addFeature(feat)
        
        layer.updateExtents()

        line_color = '#333333'
        # 设置单一符号渲染
        symbol = QgsLineSymbol.createSimple({
            'line_color': line_color,
            'line_width': 0.6,
            'line_style': 'solid'
        })
        
        layer.renderer().setSymbol(symbol)
        
        # 触发重绘
        layer.triggerRepaint()
        
        QgsProject.instance().addMapLayer(layer)
        print(f'Loaded FDEB edges from {fdeb_edges_path} into QGIS.')

